{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Viterbi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]='3'\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import sys\n",
    "sys.path.append('waveglow/')\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import IPython.display as ipd\n",
    "import pickle as pkl\n",
    "import torch\n",
    "import hparams\n",
    "from torch.utils.data import DataLoader\n",
    "from modules.model import Model\n",
    "from text import text_to_sequence, sequence_to_text\n",
    "from denoiser import Denoiser\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "import librosa\n",
    "\n",
    "data_type = 'phone'\n",
    "checkpoint_path = f\"training_log/transformer-tts-{data_type}/checkpoint_200000\"\n",
    "state_dict = {}\n",
    "for k, v in torch.load(checkpoint_path)['state_dict'].items():\n",
    "    state_dict[k[7:]]=v\n",
    "\n",
    "\n",
    "model = Model(hparams).cuda()\n",
    "model.load_state_dict(state_dict)\n",
    "_ = model.cuda().eval()\n",
    "\n",
    "datasets = ['train', 'val', 'test']\n",
    "\n",
    "for dataset in datasets:\n",
    "    with open(f'filelists/ljs_audio_text_{dataset}_filelist.txt', 'r') as f:\n",
    "        lines = [line.split('|') for line in f.read().splitlines()]\n",
    "\n",
    "    for i in tqdm(range(len(lines))):\n",
    "        file_name, _, text = lines[i]\n",
    "        text = '^' + text + '~'\n",
    "        seq = os.path.join('/media/disk1/lyh/LJSpeech-1.1/preprocessed',\n",
    "                           f'{data_type}_seq')\n",
    "        mel = os.path.join('/media/disk1/lyh/LJSpeech-1.1/preprocessed',\n",
    "                           'melspectrogram')\n",
    "\n",
    "        with open(f'{seq}/{file_name}_sequence.pkl', 'rb') as f:\n",
    "            text_padded = pkl.load(f).unsqueeze(0)\n",
    "        with open(f'{mel}/{file_name}_melspectrogram.pkl', 'rb') as f:\n",
    "            mel_padded = pkl.load(f).unsqueeze(0)\n",
    "            \n",
    "        text_lengths=torch.LongTensor([text_padded.size(1)])\n",
    "        mel_lengths=torch.LongTensor([mel_padded.size(2)])\n",
    "        \n",
    "        _, melspec, _, _, alignments, _ = model.outputs(text_padded.cuda(),\n",
    "                                                        mel_padded.cuda(),\n",
    "                                                        text_lengths.cuda(),\n",
    "                                                        mel_lengths.cuda())\n",
    "\n",
    "\n",
    "        melspec = melspec[0]\n",
    "        alignments = alignments[0]\n",
    "        \n",
    "        F = torch.mean(torch.max(alignments, dim=-1)[0], dim=-1)\n",
    "        r, c = torch.argmax(F).item()//4, torch.argmax(F).item()%4\n",
    "        location = torch.max(alignments[r,c], dim=1)[1]\n",
    "        \n",
    "        diff = location[1:]-location[:-1]\n",
    "        # First alignment easily failed\n",
    "        init_att = torch.zeros(1, alignments.size(-1))\n",
    "        init_att[0,0]=1\n",
    "        diff = diff[1:]\n",
    "        valid_data = torch.all(torch.abs(diff)<=3)\n",
    "\n",
    "        if valid_data:\n",
    "            with open(f'{hparams.teacher_path}/targets/{file_name}.pkl', 'wb') as f:\n",
    "                pkl.dump(melspec.detach().cpu(), f)\n",
    "            with open(f'{hparams.teacher_path}/alignments/{file_name}.pkl', 'wb') as f:\n",
    "                align = torch.cat([init_att, alignments[r,c, 1:].detach().cpu()], dim=0)\n",
    "                pkl.dump(align, f)\n",
    "                \n",
    "for i in range(hparams.n_layers):\n",
    "    fig, axes=plt.subplots(1,4, figsize=(20,5))\n",
    "    for j in range(hparams.n_heads):\n",
    "        axes[j].imshow(alignments[i][j].detach().cpu().t(), origin='lower', aspect='auto')\n",
    "        axes[j].set_title(f'F:{F[i,j]}')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_pytorch_p36)",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
